#!/usr/bin/env python
import os
import sys
import re
import shutil

## Usage: python rename_files.py output_dir input_dir prefix_number
## Last two arguments are optional

work_dir = os.getcwd()
prefix_number = 0
file_count=0

output_dir=os.path.abspath(sys.argv[1])

if len(sys.argv) > 2:
    work_dir=os.path.abspath(sys.argv[2])

if len(sys.argv) > 3:
    prefix_number = int(sys.argv[3])

print 'Reading from: ' + work_dir
print 'Outputting to: ' + output_dir
print 'Prefix #: ' + str(prefix_number)

seen_files=[]
new_files=[]
written_files=[]
unique_ids_to_write=[]
output_suffix=1

fileIn_path=os.path.join(output_dir,'seen_files.txt')

if os.access(fileIn_path,os.F_OK):
    fileIn = open(fileIn_path,'r')
    for line in fileIn:
        seen_files.append(line.strip('\n'))
    output_suffix=int(seen_files.pop(0))
    fileIn.close()
else:
    print 'Creating new input file'

print 'Output suffix: ' + str(output_suffix)

#print seen_files

for root, dirs, infiles in os.walk(work_dir):
    for infile in infiles:
        m = re.search('(1a2k[_])(\d*)([._])(.*?)',infile)
        if m:
            full_file = os.path.join(root,infile)
            #print full_file
            if full_file not in seen_files:
                new_files.append(full_file)

print 'New files found: ' + str(len(new_files))

file_extensions=['.ga.entities.gz','.ga.generations.gz','_backrub.out','_low.pdb.gz','_seqtol.out']

for new_file in new_files:
    if new_file[len(new_file)-11:]=='_seqtol.out':
        unique_id=new_file[:len(new_file)-11]
        match_count=0
        for file_extension in file_extensions:
            if unique_id+file_extension in new_files:
                match_count+=1
        if match_count==len(file_extensions):
            unique_ids_to_write.append(unique_id)

print 'New iterations finished: ' + str(len(unique_ids_to_write))

for unique_id_to_write in unique_ids_to_write:
    for file_extension in file_extensions:
        file_to_copy=unique_id_to_write+file_extension
        written_files.append(file_to_copy)
        destination_file = os.path.join(output_dir,'1a2k_'+str(prefix_number*1000+output_suffix)+file_extension)
        #print 'Source File: ' + file_to_copy
        #print 'Destination File: ' + destination_file
        #print ''
        shutil.copy2(file_to_copy,destination_file)
        file_count+=1
    output_suffix+=1

print 'Files copied: ' + str(file_count)

fileOut=open(os.path.join(output_dir,'seen_files.txt'),'w')
fileOut.write(str(output_suffix)+'\n')
for seen_file in seen_files:
    fileOut.write(seen_file+'\n')
for written_file in written_files:
    fileOut.write(written_file+'\n')

fileOut.close()
